import pandas as pd
import torch
from sklearn.discriminant_analysis import StandardScaler
from torch.utils.data import Dataset


# 自定义Dataset类
class MyDataSet(Dataset):
    def __init__(self, csv_file):
        # 读取数据
        # self.data = pd.read_csv(csv_file)
        # 读取数据，避免自动转换整数类型
        self.data = pd.read_csv(csv_file, float_precision='round_trip')

        # 打印前5行数据，确认是否有小数
        # print("原始数据：\n", self.data.head())

        # 提取特征（特征列为所有列，去掉标签列）
        self.features = self.data.drop(
            columns=[
                "*label_Dos",
                "*label_Probe",
                "*label_R2L",
                "*label_U2R",
                "*label_normal",
            ]
        ).astype(float)
        self.features = self.features.values

        # # 对特征进行标准化
        # self.scaler = StandardScaler()
        # self.features = self.scaler.fit_transform(self.features)
        # print('self.features:',self.features)
        # 提取标签
        self.labels = self.convert_labels(self.data)

    def __len__(self):
        # 返回数据集大小
        return len(self.data)

    def __getitem__(self, idx):
        # 获取每个样本的特征和标签
        x = torch.tensor(self.features[idx], dtype=torch.float32)
        y = torch.tensor(self.labels[idx], dtype=torch.float32)
        return x, y

    def convert_labels(self, df):
        # 将最后5列标签转化为数值类型
        labels = df[
            ["*label_Dos", "*label_Probe", "*label_R2L", "*label_U2R", "*label_normal"]
        ].values
        return labels


# 超参数设置
# 超参数设置
